from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_core.documents import Document
from langchain.chains.combine_documents import create_stuff_documents_chain
import os

os.environ['OPENAI_API_KEY'] = 'sk-ioXCJ3W2IQlpDzHu268aEeF7CcAf4b1f9f8c46A248C74a9d'
os.environ['OPENAI_BASE_URL'] = 'https://api.xiaoai.plus/v1'

if __name__ == '__main__':

    llm = ChatOpenAI()
    llm.model_name = 'gpt-4'

    prompt = ChatPromptTemplate.from_template(
        """仅根据提供的上下文回答以下问题:
        <上下文>
        {context}
        </上下文>
        问题: {input}"""
    )

    document_chain = create_stuff_documents_chain(llm,prompt)

    result = document_chain.invoke({
        "input": "数据空间研究院是谁出资建立的？",
        "context": [Document(
            page_content="合肥综合性国家科学中心数据空间研究院是由安徽省人民政府发起成立的事业单位，是新一代信息技术数据空间领域的大型新型研发机构，致力于引领网络空间安全和数据要素创新技术前沿和创新方向，凝聚一批海内外领军科学家团队，汇聚相关行业大数据，开展数据空间基础理论、体系架构、关键技术研究以及相关能力建设，打造大数据发展新高地，推进“数字江淮”建设，为数字中国建设贡献“安徽智慧”“合肥智慧”。")]
    })

    print(result)

